absl-py
clu
jax>=0.2.1
flax
mediapy
ml_collections
numpy>=1.21.5
etils>=0.6.0
tensorflow>=2.7.0
optax
dataclasses
tensorflow_datasets